# 空间权重变特征图
import torch
import torch.nn as nn
import torch.nn.functional as F
from .line_conv import LineConv2d


class DCLC(nn.Module):
    """ Directional Consistency and Line Type Convolution based Multi Scale Network """
    def __init__(self, input_channel=1):
        super(DCLC, self).__init__()

        # Sobel算子计算图像梯度
        self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3).cuda()
        self.sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3).cuda()

        self.down1 = nn.Sequential(
            nn.Conv2d(input_channel, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.downsample1 = nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0)
        self.Bottleneck1_0 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )
        self.Bottleneck1_1 = nn.Sequential(
            nn.Conv2d(128, 64, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )

        self.down2 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
        self.mp2 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
        self.Bottleneck2 = nn.Sequential(
            nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )

        self.down3 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )
        self.mp3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )
        self.Bottleneck3 = nn.Sequential(
            nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )

        self.down4 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 1024, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
        )
        self.mp4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(512, 1024, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
        )
        self.Bottleneck4 = nn.Sequential(
            nn.Conv2d(1024, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 1024, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
        )

        self.angle_act = nn.ReLU6(inplace=True)

        self.ps3 = nn.PixelShuffle(2)
        self.fuse3 = nn.Sequential(
            nn.Conv2d(768, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )

        self.ps2 = nn.PixelShuffle(2)
        self.fuse2 = nn.Sequential(
            nn.Conv2d(384 + 2, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.angle_base2_0 = nn.Sequential(
            nn.Conv2d(256, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 1, kernel_size=1, stride=1, padding=0),
        )
        self.angle_base2_1 = nn.Conv2d(2, 1, kernel_size=3, stride=1, padding=1)
        self.encode2 = nn.Sequential(
            nn.Conv2d(256, 8, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(8),
            nn.ReLU(inplace=True),
        )
        self.lc2 = LineConv2d(8, 8, kernel_size=3, stride=1, padding=4, modulation=True)
        self.decode2 = nn.Sequential(
            nn.Conv2d(8, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )

        self.ps1 = nn.PixelShuffle(2)
        self.fuse1 = nn.Sequential(
            nn.Conv2d(192 + 2, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.angle_base1_0 = nn.Sequential(
            nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(inplace=True),
            nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0),
        )
        self.angle_base1_1 = nn.Conv2d(2, 1, kernel_size=3, stride=1, padding=1)
        self.encode1 = nn.Sequential(
            nn.Conv2d(128, 4, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(4),
            nn.ReLU(inplace=True),
        )
        self.lc1 = LineConv2d(4, 4, kernel_size=3, stride=1, padding=4, modulation=True)
        self.decode1 = nn.Sequential(
            nn.Conv2d(4, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )

        self.ps0 = nn.PixelShuffle(2)
        self.fuse0 = nn.Sequential(
            nn.Conv2d(32 + 2, 32, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )
        self.angle_base0_0 = nn.Sequential(
            nn.Conv2d(32, 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(inplace=True),
            nn.Conv2d(8, 2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True),
            nn.Conv2d(2, 1, kernel_size=1, stride=1, padding=0),
        )
        self.angle_base0_1 = nn.Conv2d(2, 1, kernel_size=3, stride=1, padding=1)
        self.encode0 = nn.Sequential(
            nn.Conv2d(32, 2, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True),
        )
        self.lc0 = LineConv2d(2, 2, kernel_size=3, stride=1, padding=4, modulation=True)
        self.decode0 = nn.Sequential(
            nn.Conv2d(2, 32, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
        )



    def forward(self, x):
        gradx0 = F.conv2d(x, self.sobel_x, stride=1, padding=1)
        grady0 = F.conv2d(x, self.sobel_y, stride=1, padding=1)
        gradx1 = F.interpolate(gradx0, scale_factor=0.5, mode='bilinear', align_corners=False)
        grady1 = F.interpolate(grady0, scale_factor=0.5, mode='bilinear', align_corners=False)
        gradx2 = F.interpolate(gradx1, scale_factor=0.5, mode='bilinear', align_corners=False)
        grady2 = F.interpolate(grady1, scale_factor=0.5, mode='bilinear', align_corners=False)


        # 1/2 128
        x_1_0 = self.down1(x)
        x_1_1 = self.downsample1(x_1_0) + self.Bottleneck1_0(x_1_0)
        x_1_2 = self.Bottleneck1_1(x_1_1) + x_1_1

        # 1/4 256
        x_2_0 = self.down2(x_1_2) + self.mp2(x_1_2)
        x_2_2 = self.Bottleneck2(x_2_0) + x_2_0

        # 1/8 256
        x_3_0 = self.down3(x_2_2) + self.mp3(x_2_2)
        x_3_2 = self.Bottleneck3(x_3_0) + x_3_0

        # 1/16 256
        x_4_0 = self.down4(x_3_2) + self.mp4(x_3_2)
        x_4_2 = self.Bottleneck4(x_4_0) + x_4_0

        # 融合1/8
        x_3_3 = self.ps3(x_4_2)
        x_3_4 = self.fuse3(torch.cat((x_3_2, x_3_3), dim=1))

        # 融合1/4
        x_2_3 = self.ps2(x_3_4)
        x_2_4 = self.fuse2(torch.cat((x_2_2, x_2_3, gradx2, grady2), dim=1))
        angle2_0 = self.angle_base2_0(x_2_4)
        avg_2 = torch.mean(x_2_4, dim=1 ,keepdim=True)
        max_2,_ = torch.max(x_2_4, dim=1, keepdim=True)
        combine2 = torch.cat((avg_2, max_2), dim=1)
        angle2_1 = self.angle_base2_1(combine2)
        angle2 = self.angle_act(angle2_0 + angle2_1)
        angle2 = angle2 * 31.5 - 10.
        x_2_5 = x_2_4 + self.decode2(self.lc2(self.encode2(x_2_4), angle2.clone().detach()))

        # 融合1/2
        x_1_3 = self.ps1(x_2_5)
        x_1_4 = self.fuse1(torch.cat((x_1_2, x_1_3, gradx1, grady1), dim=1))
        angle1_0 = self.angle_base1_0(x_1_4)
        avg_1 = torch.mean(x_1_4, dim=1, keepdim=True)
        max_1,_ = torch.max(x_1_4, dim=1, keepdim=True)
        combine1 = torch.cat((avg_1, max_1), dim=1)
        angle1_1 = self.angle_base1_1(combine1)
        angle1 = self.angle_act(angle1_0 + angle1_1)
        angle1 = angle1 * 31.5 - 10.
        x_1_5 = x_1_4 + self.decode1(self.lc1(self.encode1(x_1_4), angle1.clone().detach()))

        # 融合 1
        x_0_0 = self.ps0(x_1_5)
        x_0_1 = self.fuse0(torch.cat((x_0_0, gradx0, grady0), dim=1))
        angle0_0 = self.angle_base0_0(x_0_1)
        avg_0 = torch.mean(x_0_1, dim=1, keepdim=True)
        max_0,_ = torch.max(x_0_1, dim=1, keepdim=True)
        combine0 = torch.cat((avg_0, max_0), dim=1)
        angle0_1 = self.angle_base0_1(combine0)
        angle0 = self.angle_act(angle0_0 + angle0_1)
        angle0 = angle0 * 31.5 - 10.
        x_0_2 = x_0_1 + self.decode0(self.lc0(self.encode0(x_0_1), angle0.clone().detach()))

        # 返回 1 特征用于热力图, 1/4 特征用于特征点分支和描述符分支
        # 返回 angle0 angle1 angle2用于角度损失
        return x_0_2, x_2_5, angle0, angle1, angle2